import numpy as np
import torch
from utils.ma_utilities import to_categorical

class Num2Onehot(object):
    def __init__(self, num_classes):
        self.NUM_CLASSES = num_classes
    def __call__(self, label, output_type=np.ndarray):
        '''
        int -> [0 0 int 0]
        :param label:
        :return:
        '''
        assert isinstance(label, int)
        if output_type == np.ndarray:
            return to_categorical(label, num_classes=self.NUM_CLASSES)[0]
        elif output_type == torch.Tensor:
            return torch.from_numpy(to_categorical(label, num_classes=self.NUM_CLASSES)[0]).float()
